#  !/bin/env python3
from ollama import Client

import threading
import time
from os import walk

client = Client(host='http://10.202.94.52:20232')
filepath = "./en/"   # Put the English articles in this folder
outfpath = "./cn/"   # The result will be saved in this folder

def translate():
    for (dirpath, dirnames, filenames) in walk(filepath):
        break

    for file in filenames: 
       with open(filepath + file, 'r') as f:
           text = f.read()

       print("Translating {} (length: {}) in progress ...".format(file, len(text)))

       response = client.chat(model='glm4', 
             messages=[{
                 'role': 'user',
                 'content': '将这篇tex写的文章翻译成中文, 保留原文的tex格式，不要转换成其它格式: ' + text}],
             #options={"num_ctx": 8192}
             options={"num_ctx": 32768}  # This allows us to use a large file
        )
       print('Done')

       with open(outfpath + file + '_cn', "w") as f:
            f.write(response['message']['content'])


if __name__ == "__main__":

       start_time = time.time()
       x = threading.Thread(target=translate)
       x.start()

       while True:
          if x.is_alive():
              time.sleep(2)
              print(".",  end=" ")
          else:
              break
   
       print("Finished. Time spent: %s seconds" % (time.time() - start_time))

